import torch
import torch.nn as nn
import torch.nn.functional as F
from core.update import BasicMultiUpdateBlock
from core.extractor import MultiBasicEncoder, Feature
from core.geometry import Combined_Geo_Encoding_Volume
from core.submodule import *

try:
    autocast = torch.cuda.amp.autocast
except:
    class autocast:
        def __init__(self, enabled):
            pass
        def __enter__(self):
            pass
        def __exit__(self, *args):
            pass

# 沙漏结构，3d正则化网络R，3D UNet (特征网络的多尺度特征通过CoEx相关和激励算法来引导网络)
class hourglass(nn.Module):
    def __init__(self, in_channels):
        super(hourglass, self).__init__()

        self.conv1 = nn.Sequential(
            BasicConv(in_channels, in_channels*2, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=2, dilation=1), # 3d卷积，批归一化，LeakyReLU，输入原通道，输出原通道*2，尺寸/2
            BasicConv(in_channels*2, in_channels*2, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=1, dilation=1)) # 3d卷积，批归一化，LeakyReLU，通道不变，尺寸不变

        self.conv2 = nn.Sequential(
            BasicConv(in_channels*2, in_channels*4, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=2, dilation=1), # 3d卷积，批归一化，LeakyReLU，输入原通道*2，输出原通道*4，尺寸/2
            BasicConv(in_channels*4, in_channels*4, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=1, dilation=1)) # 3d卷积，批归一化，LeakyReLU，通道不变，尺寸不变

        self.conv3 = nn.Sequential(
            BasicConv(in_channels*4, in_channels*6, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=2, dilation=1), # 3d卷积，批归一化，LeakyReLU，输入原通道*4，输出原通道*6，尺寸/2
            BasicConv(in_channels*6, in_channels*6, is_3d=True, bn=True, relu=True, kernel_size=3, padding=1, stride=1, dilation=1)) # 3d卷积，批归一化，LeakyReLU，通道不变，尺寸不变

        self.conv3_up = BasicConv(in_channels*6, in_channels*4, deconv=True, is_3d=True, bn=True, relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) # 3d转置卷积，批归一化，LeakyReLU，输入原通道*6，输出原通道*4，尺寸*2

        self.conv2_up = BasicConv(in_channels*4, in_channels*2, deconv=True, is_3d=True, bn=True, relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) # 3d转置卷积，批归一化，LeakyReLU，输入原通道*4，输出原通道*2，尺寸*2

        self.conv1_up = BasicConv(in_channels*2, 8, deconv=True, is_3d=True, bn=False, relu=False, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) # 3d转置卷积，输入原通道*2，输出8通道，尺寸*2

        self.agg_0 = nn.Sequential(
            BasicConv(in_channels*8, in_channels*4, is_3d=True, kernel_size=1, padding=0, stride=1), # 3d卷积，批归一化，LeakyReLU，输入原通道*8，输出原通道*4，尺寸不变
            BasicConv(in_channels*4, in_channels*4, is_3d=True, kernel_size=3, padding=1, stride=1), # 3d卷积，批归一化，LeakyReLU，通道不变，尺寸不变
            BasicConv(in_channels*4, in_channels*4, is_3d=True, kernel_size=3, padding=1, stride=1)) # 3d卷积，批归一化，LeakyReLU，通道不变，尺寸不变

        self.agg_1 = nn.Sequential(
            BasicConv(in_channels*4, in_channels*2, is_3d=True, kernel_size=1, padding=0, stride=1), # 3d卷积，批归一化，LeakyReLU，输入原通道*4，输出原通道*2，尺寸不变
            BasicConv(in_channels*2, in_channels*2, is_3d=True, kernel_size=3, padding=1, stride=1), # 3d卷积，批归一化，LeakyReLU，通道不变，尺寸不变
            BasicConv(in_channels*2, in_channels*2, is_3d=True, kernel_size=3, padding=1, stride=1)) # 3d卷积，批归一化，LeakyReLU，通道不变，尺寸不变

        self.feature_att_8 = FeatureAtt(in_channels*2, 64) # 输入64通道(2d)，输出原通道*2(3d)CoEx相关和激励权重
        self.feature_att_16 = FeatureAtt(in_channels*4, 192) # 输入192通道(2d)，输出原通道*4(3d)CoEx相关和激励权重
        self.feature_att_32 = FeatureAtt(in_channels*6, 160) # 输入160通道(2d)，输出原通道*6(3d)CoEx相关和激励权重
        self.feature_att_up_16 = FeatureAtt(in_channels*4, 192) # 输入192通道(2d)，输出原通道*4(3d)CoEx相关和激励权重
        self.feature_att_up_8 = FeatureAtt(in_channels*2, 64) # 输入64通道(2d)，输出原通道*2(3d)CoEx相关和激励权重

    def forward(self, x, features):
        conv1 = self.conv1(x) # 输入1/4原尺寸成本体积(8通道)，输出8*2通道，1/8原尺寸
        conv1 = self.feature_att_8(conv1, features[1]) # 输入1/8原尺寸成本体积，1/8原尺寸特征(64通道)，输出8*2(3d)CoEx相关和激励权重

        conv2 = self.conv2(conv1) # 输入1/8原尺寸体积(8*2)通道，输出8*4通道，1/16原尺寸
        conv2 = self.feature_att_16(conv2, features[2]) # 输入1/16原尺寸成本体积，1/16原尺寸特征(192通道)，输出8*4(3d)CoEx相关和激励权重

        conv3 = self.conv3(conv2) # 输入1/16原尺寸成本体积(8*4)通道，输出8*6通道，1/32原尺寸
        conv3 = self.feature_att_32(conv3, features[3]) # 输入1/32原尺寸成本体积，1/32原尺寸特征(160通道)，输出8*6(3d)CoEx相关和激励权重

        conv3_up = self.conv3_up(conv3) # 输入1/32原尺寸体积(8*6)通道，输出8*4通道，1/16原尺寸
        conv2 = torch.cat((conv3_up, conv2), dim=1) # 在通道维度连接张量，输入conv3_up和conv2，输出8*8通道，1/16原尺寸
        conv2 = self.agg_0(conv2) # 输入8*8通道，输出8*4通道，尺寸不变
        conv2 = self.feature_att_up_16(conv2, features[2]) # 输入1/16原尺寸成本体积，1/16原尺寸特征(192通道)，输出8*4(3d)CoEx相关和激励权重

        conv2_up = self.conv2_up(conv2) # 输入1/16原尺寸成本体积(8*4)通道，输出8*2通道，1/8原尺寸
        conv1 = torch.cat((conv2_up, conv1), dim=1) # 在通道维度连接张量，输入conv2_up和conv1，输出8*4通道，1/8原尺寸
        conv1 = self.agg_1(conv1) # 输入8*4通道，输出8*2通道，尺寸不变
        conv1 = self.feature_att_up_8(conv1, features[1]) # 输入1/8原尺寸成本体积，1/8原尺寸特征(64通道)，输出8*2(3d)CoEx相关和激励权重

        conv = self.conv1_up(conv1) # 输入1/8原尺寸成本体积(8*2)通道，输出8通道，1/4原尺寸

        return conv

class IGEVStereo(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.args = args

        context_dims = args.hidden_dims # GRU隐藏状态的维度；默认[128,128,128]

        self.cnet = MultiBasicEncoder(output_dim=[args.hidden_dims, context_dims], norm_fn="batch", downsample=args.n_downsample) # 特征提取器的上下文网络，提取多尺度特征，默认输入[[128,128,128],[128,128,128]]
        self.update_block = BasicMultiUpdateBlock(self.args, hidden_dims=args.hidden_dims) # 多尺度GRU更新块，输入GRU隐藏状态的维度；默认[128,128,128]

        self.context_zqr_convs = nn.ModuleList([nn.Conv2d(context_dims[i], args.hidden_dims[i]*3, 3, padding=3//2) for i in range(self.args.n_gru_layers)]) # 创建ModuleList，尺寸不变，数量是GRU的层数=3；默认3个2d卷积输入128通道，输出128*3通道

        self.feature = Feature() # 特征提取器的特征网络，提取多尺度特征

        self.stem_2 = nn.Sequential(
            BasicConv_IN(3, 32, kernel_size=3, stride=2, padding=1), # 2d卷积，实例归一化，LeakyReLU，输入3通道，输出32通道，尺寸/2
            nn.Conv2d(32, 32, 3, 1, 1, bias=False), # 输入32通道，输出32通道，尺寸不变
            nn.InstanceNorm2d(32), nn.ReLU() # 实例归一化，ReLU
            )
        self.stem_4 = nn.Sequential(
            BasicConv_IN(32, 48, kernel_size=3, stride=2, padding=1), # 2d卷积，实例归一化，LeakyReLU，输入32通道，输出48通道，尺寸/2
            nn.Conv2d(48, 48, 3, 1, 1, bias=False), # 输入48通道，输出48通道，尺寸不变
            nn.InstanceNorm2d(48), nn.ReLU() # 实例归一化，ReLU
            )

        self.spx = nn.Sequential(nn.ConvTranspose2d(2*32, 9, kernel_size=4, stride=2, padding=1)) # 2d转置卷积，输入2*32通道，输出9通道，尺寸*2
        self.spx_2 = Conv2x_IN(24, 32, True) # 残差块，对x：输入24通道，输出32通道，2d转置尺寸*2；拼接张量x,rem
        self.spx_4 = nn.Sequential(
            BasicConv_IN(96, 24, kernel_size=3, stride=1, padding=1), # 2d卷积，实例归一化，LeakyReLU，输入96通道，输出24通道，尺寸不变
            nn.Conv2d(24, 24, 3, 1, 1, bias=False), # 输入24通道，输出24通道，尺寸不变
            nn.InstanceNorm2d(24), nn.ReLU() # 实例归一化，ReLU
            )

        self.spx_gru = nn.Sequential(nn.ConvTranspose2d(2*32, 9, kernel_size=4, stride=2, padding=1)) # 2d转置卷积，输入2*32通道，输出9通道，尺寸*2
        self.spx_2_gru = Conv2x(32, 32, True) # 残差块，对x：输入32通道，输出32通道，2d转置尺寸*2；拼接张量x,rem

        self.conv = BasicConv_IN(96, 96, kernel_size=3, padding=1, stride=1) # 2d卷积，实例归一化，LeakyReLU，输入96通道，输出96通道，尺寸不变
        self.desc = nn.Conv2d(96, 96, kernel_size=1, padding=0, stride=1) # 输入96通道，输出96通道，尺寸不变

        self.corr_stem = BasicConv(8, 8, is_3d=True, kernel_size=3, stride=1, padding=1) # 3d卷积，批归一化，LeakyReLU，输入8通道，输出8通道，尺寸不变
        self.corr_feature_att = FeatureAtt(8, 96) # 输入96通道(2d)，输出8通道(3d)CoEx相关和激励权重
        self.cost_agg = hourglass(8) # 沙漏结构，3d正则化网络，3D UNet
        self.classifier = nn.Conv3d(8, 1, 3, 1, 1, bias=False) # 3d卷积，输入8通道，输出1通道，尺寸不变

    # 冻结批归一化层权重，使用预训练权重微调时，保留大数据集的统计信息
    def freeze_bn(self):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.eval()

    # 空间上采样
    def upsample_disp(self, disp, mask_feat_4, stem_2x):
        with autocast(enabled=self.args.mixed_precision): # 混合精度训练
            xspx = self.spx_2_gru(mask_feat_4, stem_2x) # 对mask_feat_4：输入32通道，输出32通道，1/2原尺寸；拼接张量mask_feat_4,stem_2x，输出32*2通道，1/2原尺寸
            spx_pred = self.spx_gru(xspx) # 输入2*32通道，输出9通道，原尺寸
            spx_pred = F.softmax(spx_pred, 1) # 上下文上采样权重，在第1维度(通道)softmax激活，输出(B,9,H,W)，原尺寸
            up_disp = context_upsample(disp*4., spx_pred).unsqueeze(1) # 上下文上采样，输入GRU更新的视差(1/4原尺寸)*4，上下文上采样权重

        return up_disp

    def forward(self, image1, image2, iters=12, flow_init=None, test_mode=False):
        """ Estimate disparity between pair of frames """

        image1 = (2 * (image1 / 255.0) - 1.0).contiguous() # 图像归一化到[-1,1]，返回连续存储的张量
        image2 = (2 * (image2 / 255.0) - 1.0).contiguous()
        with autocast(enabled=self.args.mixed_precision): # 混合精度训练
            features_left = self.feature(image1) # 特征提取器的特征网络，提取1/4、1/8、1/16、1/32原尺寸的特征，输出1/4、1/8、1/16、1/32原尺寸的多尺度特征
            features_right = self.feature(image2)
            stem_2x = self.stem_2(image1) # 输入3通道，输出32通道，1/2原尺寸
            stem_4x = self.stem_4(stem_2x) # 输入32通道，输出48通道，1/4原尺寸
            stem_2y = self.stem_2(image2)
            stem_4y = self.stem_4(stem_2y)
            features_left[0] = torch.cat((features_left[0], stem_4x), 1) # 在通道维度连接张量，输入1/4原尺寸的特征网络提取特征(48通道)和1/4原尺寸的2d卷积特征(48通道)，输出96通道，1/4原尺寸
            features_right[0] = torch.cat((features_right[0], stem_4y), 1)

            match_left = self.desc(self.conv(features_left[0])) # 输入96通道，输出96通道，1/4原尺寸
            match_right = self.desc(self.conv(features_right[0]))
            gwc_volume = build_gwc_volume(match_left, match_right, self.args.max_disp//4, 8) # 构造组相关性成本体积，输入1/4原尺寸左右特征图，1/4最大视差，分组数量，输出(B,组数,视差,H,W)，1/4原尺寸
            gwc_volume = self.corr_stem(gwc_volume) # 输入8通道(组相关性成本体积的组数=8)，输出8通道，尺寸不变
            gwc_volume = self.corr_feature_att(gwc_volume, features_left[0]) # 输入(1/4原尺寸)组相关性成本体积，1/4原尺寸特征(48通道)，输出8通道(3d)CoEx相关和激励权重
            geo_encoding_volume = self.cost_agg(gwc_volume, features_left) # 几何编码体积，3d正则化网络，输入(1/4原尺寸)组相关性成本体积，左视图多尺度特征，输出8通道，1/4原尺寸

            # Init disp from geometry encoding volume
            prob = F.softmax(self.classifier(geo_encoding_volume).squeeze(1), dim=1) # 生成概率体积，3d卷积后1/4原尺寸(B,1,最大视差,H，W)，删除为1的维度，在第1维度(视差)softmax激活，输出(B,最大视差概率,H,W)，1/4原尺寸
            init_disp = disparity_regression(prob, self.args.max_disp//4) # 回归视差，输入概率体积(1/4原尺寸)，最大视差//4，输出(B,1,H,W)，1/4原尺寸

            del prob, gwc_volume # 删除概率体积、组相关性体积

            if not test_mode:
                xspx = self.spx_4(features_left[0]) # 输入96通道，输出24通道，1/4原尺寸
                xspx = self.spx_2(xspx, stem_2x) # 对xspx：输入24通道，输出32通道，1/2原尺寸；拼接张量xspx,stem_2x，输出32*2通道，1/2原尺寸
                spx_pred = self.spx(xspx) # 输入2*32通道，输出9通道，原尺寸
                spx_pred = F.softmax(spx_pred, 1) # 上下文上采样权重，在第1维度(通道)softmax激活，张量形状(B,9,H,W)，原尺寸

            cnet_list = self.cnet(image1, num_layers=self.args.n_gru_layers) # 特征提取器的上下文网络，提取多尺度特征，输入原视图，GRU的层数；默认GRU的层数=3，输出128通道，1/4、1/8、1/16原尺寸的多尺度上下文特征
            net_list = [torch.tanh(x[0]) for x in cnet_list] # 对1/4、1/8、1/16原尺寸上下文特征，分别tanh激活，[(B,128,H,W)*3]
            inp_list = [torch.relu(x[1]) for x in cnet_list] # 对1/4、1/8、1/16原尺寸上下文特征，分别ReLU激活，[(B,128,H,W)*3]
            inp_list = [list(conv(i).split(split_size=conv.out_channels//3, dim=1)) for i,conv in zip(inp_list, self.context_zqr_convs)] # 分别对ReLU激活后的多尺度上下文特征做2d卷积，每个结果在第1维度(通道)分为3份(128通道)合为List，得到含有3个List的二维List；默认每个特征输入128通道，输出128*3通道，[[(B,128,H,W)*3]*3]

        geo_block = Combined_Geo_Encoding_Volume # 组合几何编码体积
        geo_fn = geo_block(match_left.float(), match_right.float(), geo_encoding_volume.float(), radius=self.args.corr_radius, num_levels=self.args.corr_levels) # 定义组合几何编码体积，输入1/4原尺寸左右特征图，几何编码体积，相关金字塔宽度4，相关金字塔层数2
        b, c, h, w = match_left.shape
        coords = torch.arange(w).float().to(match_left.device).reshape(1,1,w,1).repeat(b, h, 1, 1) # 生成1/4原尺寸W长度的List，改变形状为(1,1,w,1)，在第1、2维度分别重复b、h次，输出(B,H,W,1)
        disp = init_disp # 回归的视差(B,1,H,W)，1/4原尺寸
        disp_preds = []

        # GRUs iterations to update disparity
        for itr in range(iters): # ConvGRU迭代次数；默认32
            disp = disp.detach()
            geo_feat = geo_fn(disp, coords) # 组合几何编码体积，输入回归的视差(B,1,H,W)，(B,H,W,1)张量，输出组合几何编码体积(B,162,H,W)，通道(162=几何72+全对9+几何72+全对9)，1/4原尺寸，3d体积
            with autocast(enabled=self.args.mixed_precision): # 混合精度训练
                if self.args.n_gru_layers == 3 and self.args.slow_fast_gru: # Update low-res ConvGRU；默认不运行，GRU的层数=3，频繁迭代低分辨率GRU=False
                    net_list = self.update_block(net_list, inp_list, iter16=True, iter08=False, iter04=False, update=False)
                if self.args.n_gru_layers >= 2 and self.args.slow_fast_gru:# Update low-res ConvGRU and mid-res ConvGRU；默认不运行
                    net_list = self.update_block(net_list, inp_list, iter16=self.args.n_gru_layers==3, iter08=True, iter04=False, update=False)
                net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat, disp, iter16=self.args.n_gru_layers==3, iter08=self.args.n_gru_layers>=2) # 输入多尺度上下文特征[(B,128,H,W)*3]，多尺度上下文特征[[(B,128,H,W)*3]*3]，组合几何编码体积(B,162,H,W)，回归的视差(B,1,H,W)，True，True，输出当前隐藏状态，特征掩码，偏移视差

            disp = disp + delta_disp # 更新视差，原视差+偏移视差(GRU输出)
            if test_mode and itr < iters-1: # test模式只迭代1次
                continue

            # upsample predictions
            disp_up = self.upsample_disp(disp, mask_feat_4, stem_2x) # 空间上采样，输入GRU更新的视差，特征掩码，32通道1/2原尺寸左图像，输出原尺寸视差，(B,1,H,W)
            disp_preds.append(disp_up) # 上采样结果加入List，共32个

        if test_mode:
            return disp_up # test模式，返回第1次迭代空间上采样视差

        init_disp = context_upsample(init_disp*4., spx_pred.float()).unsqueeze(1) # 初始上下文上采样，输入回归的视差(1/4原尺寸)*4，上下文上采样权重，输出原尺寸视差，(B,1,H,W)
        return init_disp, disp_preds # 返回初始上下文上采样视差，32次迭代空间上采样视差List
